from PIL import Image
import random
import numpy as np
import math
import matplotlib.pyplot as plt

# 读取图像并转换为灰度图像
img = Image.open('./2.JPG').convert('L')
img1 = Image.open('./2.JPG').convert('L')
img2 = Image.open('./2.JPG').convert('L')

# 加载图片的灰度数值（二维数组形式）
img_array = img.load()

# 获取图像尺寸
width, height = img.size

# 展示灰度直方图函数


def showHistogram():
    # 创建长度为256的计数器列表
    histogram = [0] * 256

    # 遍历图像的每个像素并增加计数器
    for i in range(width):
        for j in range(height):
            pixel_value = int(img_array[i, j])
            histogram[pixel_value] += 1

    # 创建画布并绘制灰度直方图
    plt.bar(range(256), histogram)
    plt.show()


"""
图片添加椒盐噪声
实现原理： 随机地将图像中的一定比例的像素值取极大或者极小
"""


def add_salt_and_pepper_noise():

    # 设置噪声比例
    noise_ratio = 0.2
    temp = img1.load()
    # 添加椒盐噪声
    for y in range(height):
        for x in range(width):
            if random.random() < noise_ratio:
                if random.random() < 0.5:
                    temp[x, y] = 0
                else:
                    temp[x, y] = 255

    # 保存添加噪声后的图像
    img1.save('salt_and_pepper.jpg')

    # 计算信噪比
    print("椒盐噪声的信噪比为：", SNR(temp, img_array))


"""
图片中添加高斯模糊
    实现原理：每个像素点加上符合高斯分布的随机数；
        如果超过了（0-255）范围则对像素单独处理
    均值（mean）决定着高斯分布的中心位置，如果均值较大，中心向右移动；如果均值较小，那么中心向左移动。
    mean 决定了图像中的模糊的中心位置。调节mean参数图像会变亮

    方差（variance）描述着数据分布的离散程度，方差越大，高斯分布的曲线越扁平，方差越小，高斯分布的曲线越陡峭。
    var 决定了模糊的程度和强度。var越大噪声越大
"""


def add_gaussian_noise(mean=10, var=1000):
    # 计算标准差
    stddev = math.sqrt(var)
    temp = img2.load()
    # 添加高斯噪声
    for y in range(height):
        for x in range(width):
            # 生成随机数
            r = random.gauss(mean, stddev)

            # 添加噪声 每一个像素点添加上高斯分布的随机数
            value = int(temp[x, y] + r)
            if value < 0:
                value = 0
            elif value > 255:
                value = 255
            temp[x, y] = value

    # 保存添加噪声后的图像
    img2.save('gaussian.jpg')

    # 计算信噪比
    print("高斯噪声的信噪比为："+str(SNR(temp, img_array)))


# 计算信噪比
# 图像的信噪比公式：SNR = 10 * log10(Psignal / Pnoise)


def SNR(original_image, noisy_image):
    # 均方差初始化
    mse = 0
    # 计算均方差
    for x in range(width):
        for y in range(height):
            mse += (noisy_image[x, y] - original_image[x, y]) ** 2
    mse /= (img1.width * img1.height)

    # 计算原始图像的平均值
    total = 0
    avg = 0
    for x in range(width):
        for y in range(height):
            total += int(original_image[x, y])
    avg = total/(width*height)

    # SNR = 10 * log10（原始图像的像素值平均值 / MSE）。
    snr = 10 * np.log10(avg / mse)

    return snr


"""
# sobel边缘检测
其基本思想计算像素点周围像素的灰度值差分，以此来判断当前像素点是否处于边缘位置。
"""

# Sobel算子实现边缘检测


def sobel(img):
    r, c = img.shape
    new_image = np.zeros((r, c))
    new_imageX = np.zeros(img.shape)
    new_imageY = np.zeros(img.shape)
    dx = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])  # X方向
    dy = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])  # Y方向

    for i in range(r-3):
        for j in range(c-3):
            new_imageX[i+1, j+1] = (np.sum(img[i:i+3, j:j+3] * dx))**2
            new_imageY[i+1, j+1] = (np.sum(img[i:i+3, j:j+3] * dy))**2
            new_image[i+1, j+1] = (new_imageX[i+1, j+1]*new_imageX[i+1,
                                   j+1] + new_imageY[i+1, j+1]*new_imageY[i+1, j+1])**0.5
    return np.uint8(new_image)  # 无方向算子处理的图像


def edge():
    # 读取原始图像
    raw_array = np.array(img)

    # 对原始图像进行边缘检测
    raw_edge = sobel(raw_array)
    # 读取添加噪声后的图像
    gaussian_img = Image.open('./gaussian.jpg').convert('L')
    gaussian_array = np.array(gaussian_img)

    # 对添加噪声后的图像进行边缘检测
    gaussian_edge = sobel(gaussian_array)

    # 保存边缘检测结果
    Image.fromarray(raw_edge).save('raw_edge.jpg')
    Image.fromarray(gaussian_edge).save('gaussian_edge.jpg')

# 统计边缘点数


def edge_num(pic):
    count = 0
    # 读取图像
    pic_img = Image.open(pic).convert('L').load()
    for y in range(height):
        for x in range(width):
            if pic_img[x, y] == 255:
                count += 1
    return count/width/height


if __name__ == "__main__":
    print("1、图片分辨率：", width, "x", height)

    # 展示灰度直方图
    showHistogram()

    # 对图像添加噪声
    add_salt_and_pepper_noise()
    add_gaussian_noise()

    # 对图像进行边缘检测
    edge()

    # 计算边缘点数占比
    print(edge_num('raw_edge.jpg'))